import os
import cv2
import torch
import matplotlib.pyplot as plt
import numpy as np
from menet_net import MENetVNet


def test_segmentation(test_image_dir, output_dir):
    """
    使用保存的最优模型对测试图像进行分割预测并保存结果。

    参数:
    test_image_dir (str): 测试图像所在目录，图像格式为jpg。
    output_dir (str): 分割结果保存目录。
    """
    # 检查是否有可用的GPU
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # 加载最优模型
    best_model = MENetVNet().to(device)
    best_model.load_state_dict(torch.load(os.path.join(read_dir, "best_model.pth"), map_location=device, weights_only=True))
    best_model.eval()

    # 获取测试图像路径列表
    test_image_paths = sorted([os.path.join(test_image_dir, f) for f in os.listdir(test_image_dir) if f.endswith('.jpg')])

    with torch.no_grad():
        for i, image_path in enumerate(test_image_paths):
            # 加载测试图像
            img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
            img = cv2.resize(img, (512, 512)) / 255.0  # 归一化
            img_tensor = torch.tensor(img, dtype=torch.float32).unsqueeze(0).unsqueeze(0).to(device)  # 添加通道和批次维度

            # 进行分割预测
            output = best_model(img_tensor)
            output = (output > 0.5).float().cpu().numpy().squeeze()

            # 保存分割结果图像
            plt.imsave(os.path.join(output_dir, f"test_result_{i}.png"), output, cmap="gray")

if __name__ == "__main__":
    # 假设测试图像放在'test_images'目录下，分割结果同样保存在之前设置的'output_dir'目录下
    test_image_dir = "test_image"
    output_dir = "test_results"
    read_dir = "segmentation_results"
    test_segmentation(test_image_dir, output_dir)